import { Plugin, PluginKey } from "@/tiptap/pm";
import { Decoration, DecorationSet } from "@/tiptap/pm";
import { createLowlight, common } from "lowlight";

// 创建 lowlight 实例并注册常用语言
const lowlight = createLowlight(common);

// 查找文档中的所有代码块
function findCodeBlocks(doc: any) {
  const codeBlocks: any[] = [];
  doc.descendants((node: any, pos: number) => {
    if (node.type.name === "codeBlock") {
      codeBlocks.push({ node, pos });
    }
  });
  return codeBlocks;
}

// 生成语法高亮装饰
function getDecorations(doc: any) {
  const decorations: any[] = [];

  doc.descendants((node: any, pos: number) => {
    if (node.type.name !== "codeBlock") {
      return;
    }

    const language = node.attrs.language || "auto";
    const code = node.textContent;

    if (!code) {
      return;
    }

    try {
      let result;
      if (language === "auto" || !lowlight.registered(language)) {
        result = lowlight.highlightAuto(code);
      } else {
        result = lowlight.highlight(language, code);
      }

      // 将 lowlight 的结果转换为 ProseMirror decorations
      let offset = 0;
      const processNodes = (nodes: any[]) => {
        nodes.forEach((hNode: any) => {
          if (hNode.type === "text") {
            offset += hNode.value.length;
          } else if (hNode.type === "element") {
            const from = pos + 1 + offset;
            const className = hNode.properties?.className?.join(" ") || "";

            if (hNode.children) {
              const textLength = hNode.children
                .filter((child: any) => child.type === "text")
                .reduce(
                  (sum: number, child: any) => sum + child.value.length,
                  0
                );

              if (textLength > 0 && className) {
                decorations.push(
                  Decoration.inline(from, from + textLength, {
                    class: className,
                  })
                );
              }

              processNodes(hNode.children);
            }
          }
        });
      };

      processNodes(result.children);
    } catch (error) {
      console.warn("Lowlight highlighting failed:", error);
    }
  });

  return DecorationSet.create(doc, decorations);
}

// 创建 lowlight 插件
export function createLowlightPlugin() {
  return new Plugin({
    key: new PluginKey("lowlightCodeBlock"),
    state: {
      init: (_, { doc }) => {
        return getDecorations(doc);
      },
      apply: (transaction, decorationSet, oldState, newState) => {
        const oldNodeName = oldState.selection.$head.parent.type.name;
        const newNodeName = newState.selection.$head.parent.type.name;
        const oldNodes = findCodeBlocks(oldState.doc);
        const newNodes = findCodeBlocks(newState.doc);

        if (
          transaction.docChanged &&
          ([oldNodeName, newNodeName].includes("codeBlock") ||
            oldNodes.length !== newNodes.length)
        ) {
          return getDecorations(transaction.doc);
        }

        return decorationSet.map(transaction.mapping, transaction.doc);
      },
    },
    props: {
      decorations(state) {
        return this.getState(state);
      },
    },
  });
}
